// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <kllm/core/sampler.h>

namespace kllm {

    void InternalSamplerParams::setup_to_proto(SampleParams &result) const {
        result.set_seed(seed);
        result.set_n_prev(n_prev);
        result.set_n_probs(n_probs);
        result.set_min_keep(min_keep);
        result.set_top_k(top_k);
        result.set_top_p(top_p);
        result.set_min_p(min_p);
        result.set_xtc_probability(xtc_probability);
        result.set_xtc_threshold(xtc_threshold);
        result.set_typ_p(typ_p);
        result.set_temp(temp);
        result.set_dynatemp_range(dynatemp_range);
        result.set_dynatemp_exponent(dynatemp_exponent);
        result.set_penalty_last_n(penalty_last_n);
        result.set_penalty_repeat(penalty_repeat);
        result.set_penalty_freq(penalty_freq);
        result.set_penalty_present(penalty_present);
        result.set_dry_multiplier(dry_multiplier);
        result.set_dry_base(dry_base);
        result.set_dry_allowed_length(dry_allowed_length);
        result.set_dry_penalty_last_n(dry_penalty_last_n);
        result.set_mirostat(mirostat);
        result.set_mirostat_tau(mirostat_tau);
        result.set_mirostat_eta(mirostat_eta);
        result.set_penalize_nl(penalize_nl);
        result.set_ignore_eos(ignore_eos);
        result.set_no_perf(no_perf);
        result.mutable_dry_sequence_breakers()->Assign(dry_sequence_breakers.begin(), dry_sequence_breakers.end());
        result.mutable_samplers()->Assign(samplers.begin(), samplers.end());
        result.set_grammar(grammar);
        result.mutable_logit_bias()->Reserve(logit_bias.size());
        for (auto &it: logit_bias) {
            auto sit = result.mutable_logit_bias()->Add();
            sit->set_bias(it.bias);
            sit->set_token(it.token);
            sit->set_str_token("");
        }
    }

    void InternalSamplerParams::from_proto(const SampleParams &input) {
        seed = input.seed();
        n_prev = input.n_prev();
        n_probs = input.n_probs();

        min_keep = input.min_keep();
        top_k = input.top_k();
        top_p = input.top_p();
        min_p = input.min_p();
        xtc_probability = input.xtc_probability();
        xtc_threshold = input.xtc_threshold();
        typ_p = input.typ_p();
        temp = input.temp();
        dynatemp_range = input.dynatemp_range();
        dynatemp_exponent = input.dynatemp_exponent();
        penalty_last_n = input.penalty_last_n();
        penalty_repeat = input.penalty_repeat();
        penalty_freq = input.penalty_freq();
        penalty_present = input.penalty_present();
        dry_multiplier = input.dry_multiplier();
        dry_base = input.dry_base();
        dry_allowed_length = input.dry_allowed_length();
        dry_penalty_last_n = input.dry_penalty_last_n();
        mirostat = input.mirostat();
        mirostat_tau = input.mirostat_tau();
        mirostat_eta = input.mirostat_eta();
        penalize_nl = input.penalize_nl();
        ignore_eos = input.ignore_eos();
        no_perf = input.no_perf();
        grammar = input.grammar();
        dry_sequence_breakers.assign(input.dry_sequence_breakers().begin(), input.dry_sequence_breakers().end());
        for (auto it: input.samplers()) {
            samplers.push_back(static_cast<KaiSamplerType>(it));
        }

        logit_bias.reserve(input.logit_bias().size());
        for (auto &it: input.logit_bias()) {
            llama_logit_bias lo;
            lo.token = it.token();
            lo.bias = it.bias();
            logit_bias.push_back(lo);
        }

    }

    std::string InternalSamplerParams::print() const {
        char result[1024];

        snprintf(result, sizeof(result),
                 "\trepeat_last_n = %d, repeat_penalty = %.3f, frequency_penalty = %.3f, presence_penalty = %.3f\n"
                 "\tdry_multiplier = %.3f, dry_base = %.3f, dry_allowed_length = %d, dry_penalty_last_n = %d\n"
                 "\ttop_k = %d, top_p = %.3f, min_p = %.3f, xtc_probability = %.3f, xtc_threshold = %.3f, typical_p = %.3f, temp = %.3f\n"
                 "\tmirostat = %d, mirostat_lr = %.3f, mirostat_ent = %.3f",
                 penalty_last_n, penalty_repeat, penalty_freq, penalty_present,
                 dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n,
                 top_k, top_p, min_p, xtc_probability, xtc_threshold, typ_p, temp,
                 mirostat, mirostat_eta, mirostat_tau);

        return std::string(result);
    }

}  // namespace kllm
